import numpy as np
import threading


class ICModel:
    def __init__(self, graph):
        self.graph = graph
        self.nodes = list(graph.nodes())
        self.num_nodes = len(self.nodes)+5
        self.edge_dict = {}
        self._set_edge_dict()
        self.lock = threading.Lock()

    def _set_edge_dict(self):
        for node in self.nodes:
            neighbors = list(self.graph.neighbors(node))
            prob_dict = {}
            for neighbor in neighbors:
                prob_dict[neighbor] = np.random.random()
            self.edge_dict[node] = prob_dict

    def simulate(self, seeds, p=0.1, num_threads=1):
        activated = np.zeros(self.num_nodes, dtype=bool)
        for seed in seeds:
            activated[int(seed)] = True
        active_set = set(seeds)
        frontier_set = set(seeds)
        i = 0
        while len(frontier_set) > 0:
            i += 1
            next_frontier = set()
            if num_threads > 1:
                pool = []
            for node in frontier_set:
                neighbors = self.edge_dict[node]
                for neighbor in neighbors:
                    if not activated[int(neighbor)]:
                        with self.lock:
                            probability = neighbors[neighbor]
                        if np.random.random() < probability * p:
                            next_frontier.add(neighbor)
                            activated[int(neighbor)] = True
            active_set |= next_frontier
            frontier_set = next_frontier
        return active_set
